#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <stdarg.h>
#include <errno.h>
#include <pthread.h>
#include <signal.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/times.h>
#include <netinet/in.h>

#define ERROR      (-1)
#define OK         (0)
#define INVALID    (-1)
#define TRUE       (1)
#define FALSE      (0)
#define ENABLE     (1)
#define ENOERROR   (0)
#define TMERROR    ((clock_t)-1)

#define WORKERINT SIGUSR1

#define DEFAULT_PORT    (44444)
#define DEFAULT_WORKERS (0)
#define DEFAULT_BUFFER  (1024)
#define DEFAULT_MAXERR  (128)

typedef enum log_level {
  LOG_DEBUG,
  LOG_INFO,
  LOG_WARNING,
  LOG_ERROR
} log_level_t;

typedef struct thread_data {
  int self_id;
  int is_created;
  int max_fd;
  fd_set mfd;
  fd_set wfd;
  pthread_t self_th;
} thread_data_t;

typedef struct client_data {
  int size;
  char *buffer;
} client_data_t;

/*
 * Thread local storage globals
 */
__thread thread_data_t *tls_data = NULL;

/*
 * Global variables
 */
int   g_port    = DEFAULT_PORT;
int   g_debug   = FALSE;
int   g_numw    = DEFAULT_WORKERS;
int   g_bufsize = DEFAULT_BUFFER;

pthread_barrier_t *g_barrier = NULL;
pthread_mutex_t   *g_mutex   = NULL;
thread_data_t     *g_wrks    = NULL;
int                g_stop    = FALSE;
client_data_t      g_clients[FD_SETSIZE] = {{0}}; 

#define safe_free(m) do {                       \
      __safe_free(m);                           \
      m = NULL;                                 \
  } while (0)

#define log_debug(...) do {                     \
      if (g_debug) {                            \
          __log_debug(__VA_ARGS__);             \
      }                                         \
  } while (0)

#define safe_close_barrier() do {               \
      __safe_close_barrier(g_barrier);          \
      safe_free(g_barrier);                     \
  } while (0)

#define safe_close_mutex() do {                 \
      __safe_close_mutex(g_mutex);              \
      safe_free(g_mutex);                       \
  } while (0)

#define MAX(l,r)                                \
  ((l) > (r) ? (l) : (r))

/*
 * Internal functions
 */
void __safe_free(void*);
void __log_internal(log_level_t level, char *format, va_list va);
void __log_debug(char *format, ...);
void __safe_close_barrier(pthread_barrier_t *barrier);
void __safe_close_mutex(pthread_mutex_t *mutex);

/*
 * Logging utilities
 */
void __log_internal(log_level_t level, char *format, va_list va)
{
    char *formatex = NULL;
    const char *strlevel;

    switch (level) {
        case LOG_INFO:
            strlevel = "INFO";
            break;
        case LOG_WARNING:
            strlevel = "WARNING";
            break;
        case LOG_ERROR:
            strlevel = "ERROR";
            break;
        default:
            strlevel = "DEBUG";
            break;
    }
    if (asprintf(&formatex, "[%s - %p] %s\n", strlevel, tls_data, format) != INVALID) {
        vfprintf(stdout, formatex, va);
        fflush(stdout);
        safe_free(formatex);
    }
}

void __log_debug(char *format, ...)
{
    va_list va;
    int errsave;

    errsave = errno;
    va_start(va, format);
    __log_internal(LOG_DEBUG, format, va);
    va_end(va);
    errno = errsave;
}

void log_error(char *format, ...)
{
    va_list va;
    int errsave;

    errsave = errno;
    va_start(va, format);
    __log_internal(LOG_ERROR, format, va);
    va_end(va);
    errno = errsave;
}

void log_perror(char *format, ...)
{
    char buf[DEFAULT_MAXERR];
    char *formatex;
    va_list va;
    int errsave;

    errsave = errno;
    if (asprintf(&formatex, "%s (%s)", format, strerror_r(errno, buf, sizeof(buf))) != INVALID) {
        va_start(va, format);
        __log_internal(LOG_ERROR, formatex, va);
        va_end(va);
        safe_free(formatex);
    }
    errno = errsave;
}

void log_info(char *format, ...)
{
    va_list va;
    int errsave;

    errsave = errno;
    va_start(va, format);
    __log_internal(LOG_INFO, format, va);
    va_end(va);
    errno = errsave;
}

void log_warning(char *format, ...)
{
    va_list va;
    int errsave;

    errsave = errno;
    va_start(va, format);
    __log_internal(LOG_WARNING, format, va);
    va_end(va);
    errno = errsave;
}

/*
 * Usefull utilities
 */
int close_retry(int fd)
{
    errno = ENOERROR;

    while (close(fd) == INVALID && errno == EINTR) {
        errno = ENOERROR;
    }

    return errno == ENOERROR ? OK : ERROR;
}

void safe_close_socket(int sockfd)
{
    if (sockfd != INVALID && close_retry(sockfd) == ERROR) {
        log_perror("error closing socket");
    }
}

void safe_clear_set(fd_set *fdset, int size)
{
    int it;

    for (it=0; it <= size; it++) {
        if (FD_ISSET(it, fdset)) {
            safe_close_socket(it);
        }
    }
    FD_ZERO(fdset);
}

void __safe_free(void *mem)
{
    if (mem != NULL) {
        free(mem);
    }
}

void safe_free_client(int cl)
{
    safe_free(g_clients[cl].buffer);
    g_clients[cl].size = 0;
}

void safe_free_clients()
{
    int it;

    for (it=0; it < g_bufsize; it++) {
        safe_free_client(it);
    }
}

void __safe_close_barrier(pthread_barrier_t *barrier)
{
    if (barrier != NULL) {
        pthread_barrier_destroy(barrier);
    }
}

void __safe_close_mutex(pthread_mutex_t *mutex)
{
    if (mutex != NULL) {
        pthread_mutex_destroy(mutex);
    }
}

void safe_cancel_workers(thread_data_t *wrks, int size)
{
    int it;

    if (wrks) {
        for (it = 0; it < size ; it++) {
            if (wrks[it].is_created) {
                pthread_cancel(wrks[it].self_th);
                wrks[it].is_created = FALSE;
            }
        }
    }
}

void safe_stop_workers()
{
    int it;

    if (g_wrks) {
        /* Global stop flag */
        g_stop = TRUE;
        
        /* Weakup all workers to check the flag */
        for (it = 0; it < g_numw; it++) {
            pthread_kill(g_wrks[it].self_th, SIGINT);
        }

        /* Wait for the workers to terminate */
        for (it = 0; it < g_numw; it++) {
            pthread_join(g_wrks[it].self_th, NULL);
        }
    }
}

/*
 * Program routines
 */
void show_help()
{
    fprintf(stderr, 
            "Usage:\n"
            "    echo [options]'\n"
            "       -p port: listen on port\n"
            "       -d:      enable debug logs\n"
            "       -w num:  set number of workers\n");
    fflush(stderr);
}

int socket_init(int fd)
{
    int op = ENABLE;
    int errsave, res = OK;

    errsave = errno;
    if (setsockopt(fd, SOL_SOCKET,  SO_REUSEADDR, (char *)&op, sizeof(op)) == INVALID) {
        log_perror("unable to set socket options");
        goto error;
    }
    if (ioctl(fd, FIONBIO, (char *)&op) == INVALID) {
        log_perror("unable to set socket non-blocking");
        goto error;
    }

    goto end;
 error:
    res = ERROR;
 end:
    errno = errsave;

    return res;
}

int listen_socket()
{
    struct sockaddr_in saddr;
    int sockfd = INVALID;

    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) == INVALID) {
        log_error("can't create socket");
        goto error;
    }
    if (socket_init(sockfd) == ERROR) {
        goto error;
    }

    memset(&saddr, 0, sizeof(saddr));
    saddr.sin_port   = htons(g_port);
    saddr.sin_family = AF_INET;
    saddr.sin_addr.s_addr = htonl(INADDR_ANY);
    if (bind(sockfd, (struct sockaddr *) &saddr, sizeof(saddr)) == INVALID) {
        log_perror("unable to bind the socket");
        goto error;
    }
    if (listen(sockfd, 32) == INVALID) {
        log_perror("unable to set listen mode");
        goto error;
    }
    log_info("accepting connections on port %d\n", g_port);

    return sockfd;
 error:
    safe_close_socket(sockfd);
    exit(ERROR);

    return ERROR;
}

int do_accept(int sd) 
{
    errno = 0;

    return accept(sd, NULL, NULL);
}

int do_recv(int sd, char *buffer, size_t len)
{
    errno = 0;

    return recv(sd, buffer, len, 0);
}

int do_send(int sd, char *buffer, size_t len)
{
    errno = 0;

    return send(sd, buffer, len, 0);
}

int do_select(int nfds, fd_set *rfds, fd_set *wfds, fd_set *efds, struct timeval *tmout)
{
    errno = 0;

    return select(nfds, rfds, wfds, efds, tmout);
}

void server_signal_handler(int signum)
{
    switch (signum) {
        case WORKERINT:
            /* Just to interrupt the threads to pickup new clients */
            break;
        case SIGINT:
            /* This flag will stop the server */
            g_stop = TRUE;
            break;
        default:
            break;
    }
}

int setup_signals()
{
    if (signal(WORKERINT, server_signal_handler) == SIG_ERR) {
        return ERROR;
    }
    if (signal(SIGINT, server_signal_handler) == SIG_ERR) {
        return ERROR;
    }

    return OK;
}

int create_mutex()
{
    if ((g_mutex = (pthread_mutex_t*) malloc(sizeof(pthread_mutex_t))) == NULL) {
        log_error("out of memory");
        goto error;
    }
    if (pthread_mutex_init(g_mutex, NULL) != OK) {
        log_error("unable to create mutex");
        goto error;
    }

    return OK;
 error:
    safe_free(g_mutex);

    return ERROR;
}

void lock_mutex()
{
    if (pthread_mutex_lock(g_mutex) != 0) {
        log_error("failed to lock mutex");
        exit(ERROR);
    }
}

void unlock_mutex()
{
    if (pthread_mutex_unlock(g_mutex) != 0) {
        log_error("failed to unlock mutex");
        exit(ERROR);
    }
}

int create_barrier()
{
    if ((g_barrier = (pthread_barrier_t*) malloc(sizeof(pthread_barrier_t))) == NULL) {
        log_error("out of memory");
        goto error;
    }
    if (pthread_barrier_init(g_barrier, NULL, g_numw+1) != OK) {
        log_error("unable to create barrier");
        goto error;
    }

    return OK;
 error:
    safe_free(g_barrier);
    
    return ERROR;
}

void wait_barrier()
{
    int res = pthread_barrier_wait(g_barrier);

    if (res != 0 && res != PTHREAD_BARRIER_SERIAL_THREAD) {
        log_error("failed to wait for barrier");
        exit(ERROR);
    }
}

int init_client_buffer(int cid)
{
    if (cid < 0 || cid >= FD_SETSIZE) {
        log_error("invalid client");
        return ERROR;
    }
    if (g_clients[cid].buffer == NULL) {
        if ((g_clients[cid].buffer = (char*) malloc(g_bufsize)) == NULL) {
            log_error("out of memory");
            return ERROR;
        }
    }
    return OK;
}

void close_connection(int fd)
{
    safe_close_socket(fd);
    safe_free_client(fd);

    FD_CLR(fd, &tls_data->wfd);
    lock_mutex();
    FD_CLR(fd, &tls_data->mfd);
    unlock_mutex();
}

int execute_accept(int nfd, fd_set *rfd, int sockfd)
{
    if (FD_ISSET(sockfd, rfd)) {
        log_debug("begine accept new clients");
        do {
            int rc = do_accept(sockfd);

            if (rc >= 0) {
                if (socket_init(rc) == OK) {
                    thread_data_t *data = tls_data;

                    if (rc%(g_numw+1)) {
                        data = g_wrks + rc%(g_numw+1) - 1;
                    }
                    lock_mutex();
                    FD_SET(rc, &data->mfd);
                    unlock_mutex();
                    data->max_fd = MAX(data->max_fd, rc);
                    pthread_kill(data->self_th, WORKERINT);

                    log_debug("new client accepted on descriptor %d", rc);
                }
            }
        } while (errno == ENOERROR && !g_stop);
        log_debug("done accepting new clients");

        /* Clear descriptor from read set because we are done with it */
        FD_CLR(sockfd, rfd);

        nfd -= 1;
    }

    return nfd;
}

int execute_recv(int nfd, fd_set *rfd)
{
    int it;

    for (it = 0; it <= tls_data->max_fd && nfd > 0; it++) {
        if (FD_ISSET(it, rfd)) {
            int close_conn = FALSE, rc = 0;

            log_debug("receive descriptor active (%d)", it);
            if (init_client_buffer(it) == ERROR) {
                close_conn = TRUE;
            }
            else {
                do {
                    rc = do_recv(it, g_clients[it].buffer+g_clients[it].size, g_bufsize-g_clients[it].size);
                    if (rc > 0) {
                        log_debug("recv data size is %d", rc);

                        FD_SET(it, &tls_data->wfd);
                        g_clients[it].size += rc;
                    }
                    else if (errno != EWOULDBLOCK) {
                        close_conn = TRUE;
                    }
                } while (errno == ENOERROR && !close_conn && !g_stop);
            }
            if (close_conn) {
                close_connection(it);

                log_info("connection closed for client descriptor %d", it);
            }
            nfd -= 1;
        }
    }

    return nfd;
}

int execute_send(int nfd, fd_set *wfd)
{
    int it;

    for (it = 0; it <= tls_data->max_fd && nfd > 0; it++) {
        if (FD_ISSET(it, wfd)) {
            int close_conn = FALSE, rc = 0;

            log_debug("write descriptor active (%d)", it);
            do {
                log_debug("calling send with data size %d", g_clients[it].size);
                rc = do_send(it, g_clients[it].buffer+rc, g_clients[it].size);
                if (rc > 0) {
                    log_debug("sent data size is %d", rc);

                    g_clients[it].size -= rc;
                }
                else if (errno != EWOULDBLOCK) {
                    close_conn = TRUE;
                }
            } while (errno == ENOERROR && !close_conn && !g_stop && g_clients[it].size > 0);
            
            if (g_clients[it].size == 0) {
                log_debug("all data send to client %d", it);

                FD_CLR(it, &tls_data->wfd);
            }
            if (close_conn) {
                close_connection(it);

                log_info("connection closed for client descriptor %d", it);
            }
            nfd -= 1;
        }
    }

    return nfd;
}

void tls_init_data()
{
    FD_ZERO(&tls_data->mfd);
    FD_ZERO(&tls_data->wfd);

    tls_data->max_fd = 0;
}

void worker_execute()
{
    fd_set rfd, wfd;
    int nfd;

    wait_barrier();

    while (!g_stop) {
        rfd = tls_data->mfd;
        wfd = tls_data->wfd;

        log_debug("calling select from worker");
        nfd = do_select(tls_data->max_fd+1, &rfd, &wfd, NULL, NULL);

        if (nfd == INVALID) {
            log_debug("select interrupted");
            continue;
        }
        log_debug("%d descriptors are ready", nfd);

        log_debug("resolving write file descriptors");
        nfd = execute_send(nfd, &wfd);

        log_debug("resolving read file descriptors");
        nfd = execute_recv(nfd, &rfd);
    }

    safe_clear_set(&tls_data->mfd, tls_data->max_fd);
}

void *start_worker(void *data)
{
    tls_data = (thread_data_t*)data;

    tls_init_data();
    
    /* run the main routine */
    log_info("worker %d started", tls_data->self_id);
    worker_execute();
    log_info("worker %d ended", tls_data->self_id);

    return NULL;
}

int create_workers()
{
    if (g_numw > 0) {
        int it;

        if ((g_wrks = (thread_data_t*) calloc(g_numw, sizeof(thread_data_t))) == NULL) {
            log_error("out of memory");
            goto error;
        }
        for (it = 0; it < g_numw; it++) {
            if (pthread_create(&g_wrks[it].self_th, NULL, start_worker, g_wrks+it) != OK) {
                log_error("unable to create thread");
                goto error;
            }
            g_wrks[it].self_id = it;
            g_wrks[it].is_created = TRUE;
        }
    }

    return OK;
 error:
    safe_cancel_workers(g_wrks, g_numw);
    safe_free(g_wrks);

    return ERROR;
}

void main_server_execute()
{
    int sockfd, nfd;
    fd_set rfd, wfd;

    wait_barrier();

    sockfd = listen_socket();
    tls_data->max_fd = sockfd;
    FD_SET(sockfd, &tls_data->mfd);

    while (!g_stop) {
        rfd = tls_data->mfd;
        wfd = tls_data->wfd;

        log_debug("calling select");
        nfd = do_select(tls_data->max_fd+1, &rfd, &wfd, NULL, NULL);

        if (nfd == INVALID) {
            log_debug("select interrupted");
            continue;
        }
        log_debug("%d descriptors are ready", nfd);

        log_debug("resolving write file descriptors");
        nfd = execute_send(nfd, &wfd);

        log_debug("resolving accept file descriptors");
        nfd = execute_accept(nfd, &rfd, sockfd);

        log_debug("resolving read file descriptors");
        nfd = execute_recv(nfd, &rfd);
    }

    safe_clear_set(&tls_data->mfd, tls_data->max_fd);
}

int start_main_server()
{
    thread_data_t mainsrv;
    int res = OK;

    tls_data = &mainsrv;

    tls_init_data();

    if (setup_signals() == ERROR) {
        log_error("unable to set signal handler");
        goto error;
    }
    log_info("starting %d workers", g_numw);
    if (create_barrier() == ERROR) {
        goto error;
    }
    if (create_mutex() == ERROR) {
        goto error;
    }
    if (create_workers() == ERROR) {
        goto error;
    }

    log_info("main server started");
    main_server_execute();
    log_info("main server ended");

    goto cleanup;
 error:
    res = ERROR;
 cleanup:
    safe_close_barrier();
    safe_stop_workers();
    safe_free(g_wrks);
    safe_free_clients();
    safe_close_mutex();

    return res;
}

void print_times()
{
    struct tms time;

    if (times(&time) == TMERROR) {
        log_perror("error calling times");
    }
    else {
        long clk_tck = sysconf(_SC_CLK_TCK);

        log_info("echo server executed for:\n"
                 "\t%d secs (user)\n"
                 "\t%d secs (sys)\n",
                 time.tms_utime/clk_tck, time.tms_stime/clk_tck);
    }
}

int presetup_server()
{
    if (atexit(print_times)) {
        log_error("unable to set exit function");
        return ERROR;
    }

    return OK;
}

int main(int argc, char *argv[])
{
    int opt;

    while ((opt = getopt(argc, argv, "dp:w:")) != -1) {
        switch (opt) {
            case 'p':
                g_port = atoi(optarg);
                break;
            case 'd':
                g_debug = TRUE;
                break;
            case 'w':
                g_numw = atoi(optarg);
                break;
            default:
                show_help();
                return ERROR;
        }
    }
    if (presetup_server() == ERROR) {
        return ERROR;
    }
    return start_main_server();
}
